import json
import sys
import cv2
import numpy as np
import onnx_predict.onnx_utils as onnx_utils
import matplotlib.pyplot as plt
import os
import yolo_decode as yd

sys.path.append(r"D:\desktopD\aidc_local\models\keypoint_detection\evaluation_common")
import evaluation_util


def sort_keypoints(keypoints):
    """
    Sort the keypoints by their position in the image.

    @param keypoints: 4行2列的关键点坐标，numpy数组，每行代表一个关键点的坐标(x,y)
    @return: 按左上、右上、右下、左下顺序排列的关键点坐标
    """
    index_left_up = np.argmin(np.sum(keypoints * np.array([1, 1]), axis=1))
    index_left_down = np.argmin(np.sum(keypoints * np.array([1, -1]), axis=1))
    index_right_down = np.argmax(np.sum(keypoints * np.array([1, 1]), axis=1))
    index_right_up = np.argmax(np.sum(keypoints * np.array([1, -1]), axis=1))
    keypoints_sorted = np.array(
        [
            keypoints[index_left_up],
            keypoints[index_right_up],
            keypoints[index_right_down],
            keypoints[index_left_down],
        ]
    )
    return keypoints_sorted


def draw_and_save(image_name, img_rgb, out, save_folder):
    """
    绘制可视化结果并保存
    :param image_name: 图像名称,保存结果的文件名会与图像名称相同
    :param img_rgb: 用于绘制可视化结果的图像
    :param out: 神经网络输出经过后处理后的数据
    :param save_folder: 存放结果的目录
    """
    # 绘制可视化结果
    h = img_rgb.shape[0]
    colors1 = [255, 0, 0]
    # 矩形框的颜色
    colors2 = [0, 255, 255]
    # 关键点的颜色
    colors3 = [255, 255, 0]
    if out is not None:
        out = out / 256 * h
        for rect in out:
            cv2.putText(
                img_rgb,
                str(round(rect[4], 2)),
                (int(rect[0] - rect[2] // 2), int(rect[1] - rect[3] // 2)),
                1,
                5,
                colors1,
                5,
            )
            for i in [5, 8, 11, 14]:
                cv2.putText(
                    img_rgb,
                    str(round(rect[i + 2], 2)),
                    (int(rect[i] - 75), int(rect[i + 1] - 20)),
                    1,
                    2,
                    colors3,
                    2,
                )
            rect = list(map(int, rect))
            # cv2.rectangle(
            #     img_rgb,
            #     (rect[0] - rect[2] // 2, rect[1] - rect[3] // 2),
            #     (rect[0] + rect[2] // 2, rect[1] + rect[3] // 2),
            #     colors1,
            #     3,
            # )
            keypoints_array = np.array(
                [
                    [rect[5], rect[6]],
                    [rect[8], rect[9]],
                    [rect[11], rect[12]],
                    [rect[14], rect[15]],
                ]
            ).astype(np.int32)
            keypoints_array = sort_keypoints(keypoints_array)
            cv2.polylines(img_rgb, [keypoints_array], True, colors2, 3)

            r = img_rgb.shape[0] // 100
            cv2.circle(img_rgb, (int(keypoints_array[0, 0]), int(keypoints_array[0, 1])), r, colors3, -1)
            cv2.circle(img_rgb, (int(keypoints_array[1, 0]), int(keypoints_array[1, 1])), r, colors3, -1)
            cv2.circle(img_rgb, (int(keypoints_array[2, 0]), int(keypoints_array[2, 1])), r, colors3, -1)
            cv2.circle(img_rgb, (int(keypoints_array[3, 0]), int(keypoints_array[3, 1])), r, colors3, -1)

    plt.imsave(os.path.join(save_folder, image_name), img_rgb)


def main():
    img_folder = r"D:\desktopD\aidc_local\models\keypoint_detection\app_runs_result\test_barcode_img_for_pose\2D"
    json_folder = r"D:\desktopD\aidc_local\models\keypoint_detection\app_runs_result\二维码推理结果\模型1\dlc未量化"
    save_folder = r"C:\Users\Administrator\Desktop\ultralytics-action\预测结果"

    image_name_list = os.listdir(img_folder)
    for image_name in image_name_list:
        if not image_name.endswith(".png"):
            continue
        print(image_name)
        # 读取神经网络输出的json格式数据
        json_name = image_name.replace(".png", ".json")
        with open(os.path.join(json_folder, json_name), "r") as f:
            tensor1 = json.load(f)
        wait_for_decode = np.array(tensor1)

        # 读取图像，用于绘制可视化结果
        img_rgb = cv2.imread(os.path.join(img_folder, image_name))

        # 读取图像对应的标签label文件，用于计算关键点预测的指标。label文件与图像文件同名，且位于同一目录下
        label_name = image_name.replace(".png", ".json")

        # 调用decode函数，对神经网络输出进行解码，得到可以用于后处理的数据
        predicted_data = yd.yolo_decode(wait_for_decode, stride=16)

        # 调用nms_pose函数，对解码后的数据进行阈值过滤和非极大值抑制，得到最终的关键点坐标
        out = onnx_utils.nms_pose(prediction=predicted_data, conf_threshold=0.5, nms_threshold=0.6)
        out = out[0]

        # 绘制可视化结果并保存
        draw_and_save(image_name, img_rgb, out, save_folder)


if __name__ == "__main__":
    main()
